#!/usr/bin/env python
# coding=utf8

from flask import Flask, request, jsonify
from n2 import HnswIndex
import sys
import configparser
import numpy as np
import ast
from threading import *
import redis
import os
import socket
import requests


#config_center_host = 'http://172.16.237.12/'
try:
    config_url = config_center_host + 'app.ini'
    r = requests.get(config_url)
    if r.status_code == 200:
        with open("app.ini", 'wb') as f:
            f.write(r.content)
    else:
        print("loading config file error: " + config_url)
except:
    pass


class loadConfig:
    @staticmethod
    def loadConfig():
        filePath = os.path.join(os.path.dirname(os.path.abspath(__file__)) + "/app.ini")
        socketHandle = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            socketHandle.connect(('8.8.8.8', 80))
            IpAddress = socketHandle.getsockname()[0]
        except:
            print('No access to IP address')
        finally:
            socketHandle.close()

        configHandle = configparser.ConfigParser()
        configHandle.read(filePath, 'utf-8')
        groupName = configHandle.get("HOSTGROUPMAPS", IpAddress)
        sectionName = "Deploy_" + groupName
        deployOptions = [(sectionName, "Heap"), (sectionName, "Demension"), (sectionName, "IndexNoStart"),
                         (sectionName, "Depth"), (sectionName, "UseMap"), (sectionName, "AssetFamily"),
                         (sectionName, "SlbAddress"), (sectionName, "RedisConfig")]

        for sec, op in deployOptions:
            if not configHandle.has_option(sec, op):
                sys.stderr.write("ERROR: need (%s, %s) in %s.\n" % (sec, op, filePath))
                sys.stderr.write("Read README to get help inforamtion.\n")
                sys.exit(1)
        Heap = configHandle.get(sectionName, "Heap")
        Demension = configHandle.get(sectionName, "Demension")
        IndexNoStart = configHandle.get(sectionName, "IndexNoStart")
        Depth = configHandle.get(sectionName, "Depth")
        UseMap = configHandle.get(sectionName, "UseMap")
        AssetFamily = configHandle.get(sectionName, "AssetFamily")
        SlbAddress = configHandle.get(sectionName, "SlbAddress")
        RedisConfig = configHandle.get(sectionName, "RedisConfig")
        return Heap, Demension, IndexNoStart, Depth, UseMap, AssetFamily, IpAddress, SlbAddress, RedisConfig


class loadThread(Thread):
    def __init__(self, name, uServer, heap):
        Thread.__init__(self)
        self.name = name
        self.uServer = uServer
        self.heap = heap

    def run(self):
        print("loading thread %s start" % str(self.name))
        # self.result.extend(work(self.uServer, self.heap))
        loading(self.uServer, self.heap)
        print("loading thread %s end" % str(self.name))


# 多线程加载N2
def loading(uServer, heap):
    return uServer.load(heap, False)


Heap, Demension, IndexNoStart, Depth, UseMap, AssetFamily, IpAddress, SlbAddress, RedisConfig = loadConfig.loadConfig()
Heap = ast.literal_eval(Heap)
redis_config = ast.literal_eval(RedisConfig)
percent = round(100 / len(Heap))
u = {}

os.system('sync')
os.system('echo 3 > /proc/sys/vm/drop_caches')
print('Start load models......')
threads = []
for key in Heap:
    u[key] = HnswIndex(int(Demension), "angular")
    t = loadThread(key, u[key], Heap[key])
    t.start()
    threads.append(t)
for t in threads:
    t.join()
os.system('echo 3 > /proc/sys/vm/drop_caches')
print("load models 100% complete!!!!!")


class redisServer:
    def __init__(self, _redis_config):
        self.r = redis.Redis(**_redis_config)

    def getResultId(self, asset_family, index_no, index_id):
        r_key = '%d_%d_%d' % (index_no, index_id, asset_family)
        result_id = self.r.get(r_key)
        if result_id == "None":
            return False
        else:
            return result_id


class myThread(Thread):
    def __init__(self, features, uSever, inputNum, searchNum, order):
        Thread.__init__(self)
        self.result = []
        self.features = features
        self.uSever = uSever
        self.inputNum = inputNum
        self.searchNum = searchNum
        self.order = order

    def run(self):
        return self.result.extend(work(self.features, self.uSever, self.inputNum, self.searchNum, self.order))

    def get_result(self):
        try:
            return self.result  # 如果子线程不使用join方法，此处可能会报没有self.result的错误
        except Exception as e:
            print(e)
            return None


# 多线程获取N2
def work(features, uSever, inputNum, searchNum, order):
    nns = uSever.search_by_vector(features, inputNum, include_distances=True, ef_search=searchNum)
    result = []
    for i in range(int(inputNum)):
        try:
            tmp = order, nns[i][0], nns[i][1]
            result.append(tmp)
        except Exception as e:
            print(e)
            pass

    return result


# main
def doRequest(vector, inputNum):
    global u, AssetFamily, Depth, IndexNoStart
    featuresTmp = vector.split(',')
    featuresTmp = [a for a in featuresTmp if len(a) > 0]
    features = np.array(featuresTmp, dtype=np.float64)

    # 获取N2
    try:
        threads = []
        for key in u:
            t = myThread(features, u[key], int(inputNum), int(Depth), int(key))
            t.start()
            threads.append(t)

        results = []
        for t in threads:
            t.join()
            results.extend(t.get_result())

        nnsList = sorted(results, key=lambda x: x[2])

    except Exception as e:
        print(e)
        return False

    resultJson = {}
    for i in range(int(inputNum)):
        indexNo = int(nnsList[i][0]) + int(IndexNoStart)
        indexId = int(nnsList[i][1])
        score = nnsList[i][2]
        result_id = redisServer(redis_config).getResultId(int(AssetFamily), indexNo, indexId)
        resultJson[i] = result_id, score
    try:
        final_result_id_json = jsonify({'ids': resultJson, 'status': '200'})
    except Exception as e:
        print(e)
        final_result_id_json = jsonify({'ids': 'NULL', 'status': '203'})
    return final_result_id_json


app = Flask(__name__)


# start_run()

@app.route('/ann')
def compute():
    if 'vector' in request.args and 'input_k' in request.args:
        return doRequest(request.args['vector'], request.args['input_k'])
    else:
        return 'Input the correct URL..'


if __name__ == '__main__':
    # from werkzeug.contrib.fixers import ProxyFix
    # app.wsgi_app = ProxyFix(app.wsgi_app)
    app.run(host='0.0.0.0', threaded=True)